import numpy as np
from typing import List, Tuple, Set
import heapq

class PathPlanner:
    def __init__(self, grid_size: int = 200):
        self.grid_size = grid_size
        self.obstacles: Set[Tuple[int, int]] = set()
        self.grid_map = None  # 存储栅格地图
        self.grid_rows = 0
        self.grid_cols = 0
        
    def add_obstacle(self, x: int, y: int):
        """添加障碍物到网格"""
        self.obstacles.add((x, y))
    
    def load_grid_map(self, grid_map: np.ndarray):
        """
        从栅格地图加载障碍物
        
        参数:
        grid_map: 2D数组，其中0表示可行区域，1表示障碍物
        """
        self.grid_map = grid_map
        self.grid_rows, self.grid_cols = grid_map.shape
        
        # 清除现有障碍物
        self.obstacles.clear()
        
        # 加载障碍物
        for i in range(self.grid_rows):
            for j in range(self.grid_cols):
                if grid_map[i, j] == 1:  # 1表示障碍物
                    self.add_obstacle(j, i)  # 注意：(j,i)对应(x,y)
    
    def is_valid_position(self, x: int, y: int) -> bool:
        """检查位置是否有效（不是障碍物且在地图范围内）"""
        # 如果使用了栅格地图，检查坐标是否在地图范围内
        if self.grid_map is not None:
            if x < 0 or y < 0 or x >= self.grid_cols or y >= self.grid_rows:
                return False
            return self.grid_map[y, x] == 0  # 0表示可行区域
        
        # 如果没有栅格地图，仅检查是否是障碍物
        return (x, y) not in self.obstacles
    
    def get_neighbors(self, pos: Tuple[int, int]) -> List[Tuple[int, int]]:
        """获取相邻的有效位置（上下左右）"""
        x, y = pos
        neighbors = []
        # 上下左右移动
        for dx, dy in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
            new_x, new_y = x + dx, y + dy
            if self.is_valid_position(new_x, new_y):
                neighbors.append((new_x, new_y))
        return neighbors
    
    def manhattan_distance(self, start: Tuple[int, int], goal: Tuple[int, int]) -> int:
        """计算曼哈顿距离作为启发式函数"""
        return abs(start[0] - goal[0]) + abs(start[1] - goal[1])
    
    def find_path(self, start: Tuple[int, int], goal: Tuple[int, int]) -> List[Tuple[int, int]]:
        """使用A*算法寻找最短路径"""
        # 检查起点和终点是否有效
        if not self.is_valid_position(start[0], start[1]):
            raise ValueError(f"起点 {start} 不是有效位置")
        if not self.is_valid_position(goal[0], goal[1]):
            raise ValueError(f"终点 {goal} 不是有效位置")
        
        frontier = []
        heapq.heappush(frontier, (0, start))
        came_from = {start: None}
        cost_so_far = {start: 0}
        
        while frontier:
            current = heapq.heappop(frontier)[1]
            
            if current == goal:
                break
                
            for next_pos in self.get_neighbors(current):
                new_cost = cost_so_far[current] + 1
                
                if next_pos not in cost_so_far or new_cost < cost_so_far[next_pos]:
                    cost_so_far[next_pos] = new_cost
                    priority = new_cost + self.manhattan_distance(next_pos, goal)
                    heapq.heappush(frontier, (priority, next_pos))
                    came_from[next_pos] = current
        
        # 检查是否找到路径
        if goal not in came_from:
            return []  # 没有找到路径
        
        # 重建路径
        path = []
        current = goal
        while current is not None:
            path.append(current)
            current = came_from.get(current)
        
        return list(reversed(path))
    
    def plot_grid(self, path: List[Tuple[int, int]] = None, plt_axis=None, title='机器人底座移动路径规划'):
        """
        绘制栅格地图
        
        参数:
        path: 路径点列表
        plt_axis: matplotlib 轴对象，如果提供则在此轴上绘图
        title: 图表标题
        """
        import matplotlib.pyplot as plt
        import matplotlib.patches as patches
        
        # 如果没有提供 plt_axis，则创建新的图表
        if plt_axis is None:
            plt.figure(figsize=(10, 10))
            plt_axis = plt.gca()
        
        # 如果有栅格地图，绘制完整地图
        if self.grid_map is not None:
            # 创建一个空白背景
            plt_axis.set_xlim(-0.5, self.grid_cols - 0.5)
            plt_axis.set_ylim(self.grid_rows - 0.5, -0.5)  # 反转y轴使得原点在左上
            
            # 绘制栅格
            for y in range(self.grid_rows):
                for x in range(self.grid_cols):
                    cell_color = 'white'  # 可行区域
                    if self.grid_map[y, x] == 1:
                        cell_color = 'lightgray'  # 障碍物
                    rect = patches.Rectangle(
                        (x - 0.5, y - 0.5), 1, 1, 
                        linewidth=1, edgecolor='black',
                        facecolor=cell_color
                    )
                    plt_axis.add_patch(rect)
                    
                    # 添加坐标标签
                    plt_axis.text(x, y, f"({x},{y})", 
                                 ha='center', va='center', 
                                 fontsize=8, color='black')
            
            # 添加注释
            plt_axis.text(self.grid_cols/2, -0.9, 
                         "注: 0表示没有障碍物，即可行区域，1表示有障碍物。Start表示机器臂初始位置，End表示目标点位置，每格的大小为200mm*200mm",
                         ha='center', fontsize=9)
        else:
            # 如果没有栅格地图，则使用之前的方式绘制障碍物
            obstacles_x = [x for x, _ in self.obstacles]
            obstacles_y = [y for _, y in self.obstacles]
            plt_axis.scatter(obstacles_x, obstacles_y, c='red', marker='s', label='障碍物')
        
        # 绘制路径
        if path:
            path_x = [x for x, _ in path]
            path_y = [y for _, y in path]
            plt_axis.plot(path_x, path_y, 'b-', linewidth=2, label='路径')
            plt_axis.scatter([path_x[0]], [path_y[0]], c='green', marker='o', s=100, label='起点(Start)')
            plt_axis.scatter([path_x[-1]], [path_y[-1]], c='yellow', marker='*', s=100, label='终点(End)')
            
            # 添加箭头表示路径方向
            if len(path) > 1:
                for i in range(len(path) - 1):
                    x1, y1 = path[i]
                    x2, y2 = path[i + 1]
                    dx, dy = x2 - x1, y2 - y1
                    plt_axis.arrow(x1, y1, dx * 0.6, dy * 0.6, 
                                  head_width=0.2, head_length=0.2, 
                                  fc='blue', ec='blue')
        
        plt_axis.grid(True)
        plt_axis.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=3)
        plt_axis.set_title(title)
        plt_axis.set_xlabel('X坐标 (单位格)')
        plt_axis.set_ylabel('Y坐标 (单位格)')
        
        # 如果没有提供 plt_axis，则显示图表
        if plt_axis is plt.gca():
            plt.tight_layout()
            plt.show() 